#!/usr/bin/python
#coding:utf-8

# Copyright 2012 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, extract the edgels, and find
## the orientation using SQP.
##

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0][-7:] == '-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

from numpy import dot

from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors

from quaternion import Quat, random_quaternion

import simplejson

import Image

set_printoptions(precision=3)


def create_pic(job_params, framenum):
    filename = job_params['root_directory']+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
        job_params["gaussian_smoothing_factor"] > 0):
        for c in range(3):
            imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], double(job_params["gaussian_smoothing_factor"]))


    ## Intrinsic parameters
    p_point = job_params['principal_point']
    focal_distance = job_params['focal_distance']
    distortion = job_params['distortion_coefficient']
    pic = PicturePinhole(imr, focal_distance, p_point[0], p_point[1], distortion)

    return pic

if __name__ == '__main__':

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0][-7:] == '-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load images and initialize pic objects

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    framenum = int(sys.argv[2])

    pic = create_pic(job_params, framenum)

    gmethod = job_params['edge_detection_filter']
    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold']
    gfil_t = job_params['edgel_filter_threshold']
    gfil_d = job_params['edgel_filter_distance']

    pic.extract_edgels(gspc, glim, method=gmethod)

    #######
    ## Filter the edgels to remove the lonely ones.
    ## Create kd-tree with all edgel positions    
    if gfil_t > 0:
        kdt = scipy.spatial.cKDTree( pic.edgels[:,:2] )
        ok_edgels = []
        ddh = []
        for N in range(pic.edgels.shape[0]):
            px,py,ux,uy = pic.edgels[N]

            dist, idx = kdt.query(array([[px,py]]), 10, eps=0, p=2, distance_upper_bound=gfil_d)
            iidx = idx[dist<inf][1:]

            Nok = 0
            for nx, ny, vx, vy in pic.edgels[iidx]:
                dd = np.abs((nx-px) * ux + (ny-py) * uy)
                ddh.append(dd)
                if dd < gfil_t:
                    ddr = np.abs((px-nx) * vx + (py-ny) * vy)
                    if ddr < gfil_t:
                        Nok += 1
            if Nok >=2:
                ok_edgels.append(N)
        pic.edgels = pic.edgels[ok_edgels]
    ##
    #################################################################

    ###### On to calculate normals
    ##
    pic.generate_random_edgels()
    # pic.calculate_edgel_normals()
    # pic.calculate_random_normals()
    pic.calculate_edgel_normals_harris()
    pic.calculate_random_normals_harris()

    ## Run optimization
    sig = job_params['angular_threshold']

    #xini=[random_quaternion().q for x in range(100)] ## Lots of random initial points
    xini = array(job_params['initial_guesses'])   ## Read initial points form job file
    #xini = array([[1.,0,0,0]]) ## Just "no rotation"
    #xini = []
    
    qopt, sig_hat = pic.find_orientation(xini, sig, method=0)
    #xopt, trash = pic.find_orientation(xini, sig, method=1)
    sig_hat = sig
    ##
    #################################################################
    pic.orientation = qopt.canonical()

    # print 'fd = ', pic.fd
    # print 'kappa = ', pic.distortion
    print framenum,
    print qopt.q[0], qopt.q[1], qopt.q[2], qopt.q[3],
    print pic.edgels.shape[0]
    #print pic.middlex, pic.middley, pic.fd, pic.distortion

    if PlotStuff:
        pic.set_labels(sig_hat)
        figure(0)
        plot_aux.plot_problem(pic)

        figure(1)
        plot_aux.plot_spheres(pic.orientation.rot().T, pic.normals)


    if PlotStuff:
        figure(11)
        imshow(.1+pic.frame*.9/255.)
        pic.plot_vdirs(gca(), 17, pic.orientation.rot().T)
        axis('equal')
        grid()
